import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report, confusion_matrix, precision_recall_curve, roc_curve, auc
from imblearn.over_sampling import SMOTE
import matplotlib.pyplot as plt
import seaborn as sns

# 定义数据集的列名
column_names = [
    'LOC_BLANK', 'BRANCH_COUNT', 'CALL_PAIRS', 'LOC_CODE_AND_COMMENT',
    'LOC_COMMENTS', 'CONDITION_COUNT', 'CYCLOMATIC_COMPLEXITY',
    'CYCLOMATIC_DENSITY', 'DECISION_COUNT', 'DECISION_DENSITY',
    'DESIGN_COMPLEXITY', 'DESIGN_DENSITY', 'EDGE_COUNT',
    'ESSENTIAL_COMPLEXITY', 'ESSENTIAL_DENSITY', 'LOC_EXECUTABLE',
    'PARAMETER_COUNT', 'GLOBAL_DATA_COMPLEXITY', 'GLOBAL_DATA_DENSITY',
    'HALSTEAD_CONTENT', 'HALSTEAD_DIFFICULTY', 'HALSTEAD_EFFORT',
    'HALSTEAD_ERROR_EST', 'HALSTEAD_LENGTH', 'HALSTEAD_LEVEL',
    'HALSTEAD_PROG_TIME', 'HALSTEAD_VOLUME', 'MAINTENANCE_SEVERITY',
    'MODIFIED_CONDITION_COUNT', 'MULTIPLE_CONDITION_COUNT', 'NODE_COUNT',
    'NORMALIZED_CYLOMATIC_COMPLEXITY', 'NUM_OPERANDS', 'NUM_OPERATORS',
    'NUM_UNIQUE_OPERANDS', 'NUM_UNIQUE_OPERATORS', 'NUMBER_OF_LINES',
    'PATHOLOGICAL_COMPLEXITY', 'PERCENT_COMMENTS', 'LOC_TOTAL', 'Defective'
]


# 加载数据
def load_data(filepath):
    data = pd.read_csv(filepath, names=column_names, na_values=['?'])
    # 将标签列从'Y', 'N'转换为1, 0
    data['Defective'] = data['Defective'].map({'Y': 1, 'N': 0})

    # 尝试将所有数值列转换为浮点数
    for col in data.columns[:-1]:
        try:
            data[col] = pd.to_numeric(data[col], errors='coerce')
        except Exception as e:
            print(f"Error converting column {col}: {e}")

    # 检查是否有非数字类型的特征
    non_numeric_features = [col for col in data.columns[:-1] if not np.issubdtype(data[col].dtype, np.number)]
    if non_numeric_features:
        raise ValueError(f"Non-numeric columns found after conversion: {non_numeric_features}")

    # 处理缺失值，这里简单地删除含有缺失值的行
    data = data.dropna()
    features = data.iloc[:, :-1].values
    labels = data.iloc[:, -1].values

    # 打印特征和标签的前几行以确认数据格式
    print("Features:")
    print(features[:5])
    print("Labels:")
    print(labels[:5])

    return features, labels


# 数据预处理
def preprocess_data(features):
    # 检查特征是否全部为数值类型
    if not np.issubdtype(features.dtype, np.number):
        raise ValueError("Features contain non-numeric values.")

    # 特征缩放
    scaler = StandardScaler()
    features = scaler.fit_transform(features)
    return features


# 过采样
def oversample_data(features, labels):
    smote = SMOTE(random_state=42)
    X_resampled, y_resampled = smote.fit_resample(features, labels)
    return X_resampled, y_resampled


# 训练KNN分类器
def train_KNN(train_data, train_labels):
    knn = KNeighborsClassifier(n_neighbors=5)  # 选择k=5作为示例
    knn.fit(train_data, train_labels)
    return knn


# 测试分类器
def testing_KNN(knn, test_data, test_labels):
    predictions = knn.predict(test_data)
    return predictions


# 绘制混淆矩阵热力图
def plot_confusion_matrix(y_true, y_pred):
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.title('Confusion Matrix')
    plt.show()


# 绘制精确率-召回率曲线
def plot_precision_recall_curve(y_true, y_scores):
    precision, recall, _ = precision_recall_curve(y_true, y_scores)
    plt.figure(figsize=(8, 6))
    plt.plot(recall, precision, marker='.')
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.title('Precision-Recall Curve')
    plt.show()


# 绘制ROC曲线
def plot_roc_curve(y_true, y_scores):
    fpr, tpr, _ = roc_curve(y_true, y_scores)
    roc_auc = auc(fpr, tpr)
    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()


def train_model(file_path):
    features, labels = load_data(file_path)
    features = preprocess_data(features)
    # 过采样处理数据不平衡问题
    features, labels = oversample_data(features, labels)

    # 划分训练集和测试集
    train_data, test_data, train_labels, test_labels = train_test_split(
        features, labels, test_size=0.2, random_state=42)

    # 训练模型
    knn = train_KNN(train_data, train_labels)
    return knn, train_data, train_labels, test_data, test_labels, features, labels


def show_prediction_results(predictions, test_labels, knn, test_data):
    # 评估模型
    print("Classification Report:")
    print(classification_report(test_labels, predictions))
    print("\nConfusion Matrix:")
    print(confusion_matrix(test_labels, predictions))

    # 绘制混淆矩阵热力图
    plot_confusion_matrix(test_labels, predictions)

    # 获取预测概率
    y_scores = knn.predict_proba(test_data)[:, 1]

    # 绘制精确率-召回率曲线
    plot_precision_recall_curve(test_labels, y_scores)

    # 绘制ROC曲线
    plot_roc_curve(test_labels, y_scores)


# 主函数
def main():
    filepath = r'C:\Users\xR\Desktop\python\software-defect-prediction\dataset\PC1.arff'  # 替换为你的数据集路径

    # 训练模型
    knn, train_data, train_labels, test_data, test_labels, features, labels = train_model(filepath)

    # 测试模型
    predictions = testing_KNN(knn, test_data, test_labels)

    # 评估模型
    show_prediction_results(predictions, test_labels, knn, test_data)


if __name__ == '__main__':
    main()
